{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/mlop-ai/mlop/blob/main/examples/torch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "<h1 align=\"center\" style=\"font-family: Inter, sans-serif; font-style: normal; font-weight: 700; font-size: 72px\">m:lop</h1>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -Uq \"mlop[full]\"\n",
    "# %pip install \"mlop[full] @ git+https://github.com/mlop-ai/mlop.git\"\n",
    "# import sys; import os; sys.path.insert(0, os.path.dirname(os.path.abspath(os.path.dirname(\"__file__\"))))\n",
    "import mlop\n",
    "\n",
    "mlop.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the experiment run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "\n",
    "torch.backends.cudnn.deterministic = True\n",
    "random.seed(hash(\"this\") % 2**32 - 1)\n",
    "np.random.seed(hash(\"improves\") % 2**32 - 1)\n",
    "torch.manual_seed(hash(\"reproducibility\") % 2**32 - 1)\n",
    "torch.cuda.manual_seed_all(hash(\"repeatability\") % 2**32 - 1)\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = dict(\n",
    "    epochs=5,\n",
    "    classes=10,\n",
    "    kernels=[16, 32],\n",
    "    batch_size=128,\n",
    "    learning_rate=0.005,\n",
    "    dataset=\"MNIST\",\n",
    "    architecture=\"CNN\",\n",
    ")\n",
    "\n",
    "\n",
    "def get_data(slice=5, train=True):\n",
    "    full_dataset = torchvision.datasets.MNIST(\n",
    "        root=\".\", train=train, transform=transforms.ToTensor(), download=True\n",
    "    )\n",
    "    sub_dataset = torch.utils.data.Subset(\n",
    "        full_dataset, indices=range(0, len(full_dataset), slice)\n",
    "    )\n",
    "    return sub_dataset\n",
    "\n",
    "\n",
    "def make_loader(dataset, batch_size):\n",
    "    loader = torch.utils.data.DataLoader(\n",
    "        dataset=dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=True,\n",
    "        pin_memory=True,\n",
    "        num_workers=8,\n",
    "    )\n",
    "    return loader\n",
    "\n",
    "\n",
    "def make(config):\n",
    "    train, test = get_data(train=True), get_data(train=False)\n",
    "    train_loader = make_loader(train, batch_size=config[\"batch_size\"])\n",
    "    test_loader = make_loader(test, batch_size=config[\"batch_size\"])\n",
    "\n",
    "    model = ConvNet(config[\"kernels\"], config[\"classes\"]).to(device)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=config[\"learning_rate\"])\n",
    "\n",
    "    return model, train_loader, test_loader, criterion, optimizer\n",
    "\n",
    "\n",
    "def train(model, loader, criterion, optimizer, config):\n",
    "    mlop.watch(model, disable_graph=True, freq=10)\n",
    "\n",
    "    total_batches = len(loader) * config[\"epochs\"]\n",
    "    print(f\"Total batches: {total_batches}\")\n",
    "    example_ct = 0  # number of examples seen\n",
    "    batch_ct = 0\n",
    "    for epoch in tqdm(range(config[\"epochs\"])):\n",
    "        for _, (images, labels) in enumerate(loader):\n",
    "            loss = train_batch(images, labels, model, optimizer, criterion)\n",
    "            example_ct += len(images)\n",
    "            batch_ct += 1\n",
    "\n",
    "            if ((batch_ct + 1) % 25) == 0:\n",
    "                train_log(loss, example_ct, epoch)\n",
    "\n",
    "\n",
    "def train_batch(images, labels, model, optimizer, criterion):\n",
    "    images, labels = images.to(device), labels.to(device)\n",
    "    outputs = model(images)\n",
    "    loss = criterion(outputs, labels)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "def train_log(loss, example_ct, epoch):\n",
    "    loss = float(loss)\n",
    "    print(f\"Loss after {str(example_ct).zfill(5)}\" + f\" examples: {loss:.3f}\")\n",
    "\n",
    "\n",
    "def test(model, test_loader):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        correct, total = 0, 0\n",
    "        for images, labels in test_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "        print(\n",
    "            f\"Accuracy of the model on the {total} \"\n",
    "            + f\"test images: {correct / total:%}\"\n",
    "        )\n",
    "    torch.onnx.export(model, images, \"model.onnx\")\n",
    "\n",
    "\n",
    "def model_pipeline(config):\n",
    "    op = mlop.init(project=\"pytorch-demo\", config=config)\n",
    "    model, train_loader, test_loader, criterion, optimizer = make(config)\n",
    "    print(model)\n",
    "\n",
    "    train(model, train_loader, criterion, optimizer, config)\n",
    "    test(model, test_loader)\n",
    "    op.finish()\n",
    "    return model\n",
    "\n",
    "\n",
    "class ConvNet(nn.Module):\n",
    "    def __init__(self, kernels, classes=10):\n",
    "        super(ConvNet, self).__init__()\n",
    "\n",
    "        self.layer1 = nn.Sequential(\n",
    "            nn.Conv2d(1, kernels[0], kernel_size=5, stride=1, padding=2),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        )\n",
    "        self.layer2 = nn.Sequential(\n",
    "            nn.Conv2d(16, kernels[1], kernel_size=5, stride=1, padding=2),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        )\n",
    "        self.fc = nn.Linear(7 * 7 * kernels[-1], classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.layer1(x)\n",
    "        out = self.layer2(out)\n",
    "        out = out.reshape(out.size(0), -1)\n",
    "        out = self.fc(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start training with **mlop**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model_pipeline(config)\n",
    "\n",
    "mlop.finish()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
